import socket
import threading
import time
import os
import json
from tkinter import Tk, scrolledtext, Frame, Button, Label, Entry, StringVar, messagebox, ttk, filedialog
import hashlib


def get_local_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(("8.8.8.8", 80))
    local_ip = s.getsockname()[0]
    s.close()
    return local_ip


class ChatServer:
    def __init__(self):
        self.window = Tk()
        self.window.title("COLServer")
        self.window.geometry("800x600")
        self.window.resizable(False, False)
        self.window.protocol("WM_DELETE_WINDOW", self.on_closing)
        # 服务器状态
        self.is_running = False
        self.clients = {}
        self.file_transfers = {}
        self.file_server_socket = None
        self.file_port = 12346
        self.lock = threading.Lock()
        self.HEARTBEAT_INTERVAL = 30
        self.HEARTBEAT_TIMEOUT = 45
        self._init_heartbeat_check()
        self._create_widgets()
        self._setup_style()

    def _init_heartbeat_check(self):
        """初始化心跳检测线程"""
        self.heartbeat_thread = threading.Thread(
            target=self._check_heartbeats,
            daemon=True
        )
        self.heartbeat_thread.start()

    def _check_heartbeats(self):
        """定期检查客户端心跳"""
        while self.is_running:
            time.sleep(self.HEARTBEAT_INTERVAL)
            current_time = time.time()

            with self.lock:  # 加锁获取客户端副本
                # 创建副本避免遍历时修改字典
                clients_copy = self.clients.copy()

            for client, info in clients_copy.items():
                last_active = info.get('last_active', 0)
                if current_time - last_active > self.HEARTBEAT_TIMEOUT:
                    self._log(f"客户端 {info['nickname']} 心跳超时，断开连接")
                    self._cleanup_client(client)

    def _create_widgets(self):
        """创建界面组件"""
        # 控制面板
        control_frame = Frame(self.window)
        control_frame.pack(pady=10, fill="x")

        # 网络配置
        Label(control_frame, text="IP地址:").pack(side="left", padx=5)
        ipaddr = get_local_ip()
        self.ip_var = StringVar(value=ipaddr)
        ttk.Entry(control_frame, textvariable=self.ip_var, width=15, state="readonly").pack(side="left")
        Label(control_frame, text="端口:").pack(side="left", padx=5)
        self.port_var = StringVar(value="12345")
        ttk.Entry(control_frame, textvariable=self.port_var, width=8).pack(side="left")
        # 服务器控制按钮
        self.btn_start = ttk.Button(control_frame, text="启动服务器", command=self.toggle_server)
        self.btn_start.pack(side="left", padx=5)

        Label(control_frame, text="     在线用户:                                   ").pack(side="right", padx=5)
        # 导出日志
        export_log_btn = ttk.Button(control_frame, text="导出日志", command=self.export_log)
        export_log_btn.pack(side="right", padx=1)
        # 清空日志
        clear_log_btn = ttk.Button(control_frame, text="清空日志", command=self.log_area_delete)
        clear_log_btn.pack(side="right", padx=1)

        # 在线用户列表
        self.user_list = scrolledtext.ScrolledText(
            self.window,
            width=25,
            height=20,
            state="disabled"
        )
        self.user_list.pack(side="right", fill="y", padx=5)

        # 日志显示区域
        self.log_area = scrolledtext.ScrolledText(
            self.window,
            wrap="word",
            state="disabled"
        )
        self.log_area.pack(expand=True, fill="both", padx=5)

        # 底部状态栏
        self.status_label = Label(self.window, text="Server Status: Stopped", anchor="w")
        self.status_label.pack(side="bottom", fill="x")

    def _setup_style(self):
        """配置界面样式"""
        self.window.configure(bg="#F0F0F0")
        self.log_area.configure(bg="white", font=("Consolas", 10))
        self.user_list.configure(bg="#F8F8F8", font=("Arial", 9))
        self.status_label.configure(bg="#E0E0E0", font=("Arial", 9))

    def toggle_server(self):
        """启动/停止服务器"""
        if not self.is_running:
            ip = self.ip_var.get()
            port = self.port_var.get()
            if not port.isdigit():
                messagebox.showerror("错误", "端口号必须为数字")
                return
            self.start_server(ip, int(port))
        else:
            self.stop_server()

    def start_server(self, ip, port):
        """启动服务器"""
        try:
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.bind((ip, port))
            self.server_socket.listen(5)
            self.is_running = True
            self.btn_start.config(text="停止服务器")
            self.status_label.config(text=f"Server Status: Running (host: {ip}, port: {port})")
            threading.Thread(target=self._accept_clients, daemon=True).start()
            self._log(f"聊天服务端口启动成功 [{ip}:{port}]")
            # 新增：启动文件传输服务器
            self.start_file_server(ip, self.file_port)
        except Exception as e:
            self._log(f"服务器启动失败: {str(e)}", "ERROR")

    def start_file_server(self, ip, port):
        """启动文件传输服务器"""
        try:
            self.file_server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.file_server_socket.bind((ip, port))
            self.file_server_socket.listen(2)
            self._log(f"文件传输端口启动成功 [{ip}:{port}]")
            threading.Thread(target=self._accept_file_clients, daemon=True).start()
        except Exception as e:
            self._log(f"文件传输端口启动失败: {str(e)}", "ERROR")

    def stop_server(self):
        """停止服务器"""
        self.is_running = False
        try:
            for client in self.clients.copy():
                client.close()
            self.server_socket.close()
            # 新增：关闭文件传输socket
            if self.file_server_socket:
                self.file_server_socket.close()
        except:
            pass
        self.btn_start.config(text="启动服务器")
        self.status_label.config(text="Server Status: Stopped")
        self._log("服务器已停止")

    def _accept_clients(self):
        """接受客户端连接"""
        while self.is_running:
            try:
                client_socket, addr = self.server_socket.accept()
                threading.Thread(
                    target=self._handle_client,
                    args=(client_socket,),
                    daemon=True
                ).start()
                self._log(f"新连接来自: [{addr[0]}:{addr[1]}]")
            except:
                break

    def _accept_file_clients(self):
        """接受文件客户端连接"""
        while self.is_running:
            if self.file_server_socket is None:
                break
            try:
                file_client_socket, addr = self.file_server_socket.accept()
                threading.Thread(
                    target=self._handle_file_client,
                    args=(file_client_socket,),
                    daemon=True
                ).start()
                self._log(f"新文件连接来自: [{addr[0]}:{addr[1]}]")
            except:
                break

    def _broadcast(self, message, sender_socket=None):
        """广播消息到所有在线主端口客户端"""
        for client in self.clients.copy():
            if sender_socket is None or client != sender_socket:
                try:
                    client.send(message)
                except:
                    self._cleanup_client(client)

    def _handle_client(self, client_socket):
        """处理主端口客户端连接（仅聊天和心跳）"""
        try:
            # 获取客户端信息
            client_ip, client_port = client_socket.getpeername()
            nickname = client_socket.recv(1024).decode().strip()
            client_socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4096)
            with self.lock:
                self.clients[client_socket] = {
                    "nickname": nickname,
                    "ip": f"{client_ip}:{client_port}",
                    "last_active": time.time()
                }
                self._update_user_list()
                self._log(f"用户上线: '{nickname}'  [{client_ip}:{client_port}]", "Login")

            # 持续接收消息
            while self.is_running:
                try:
                    data = client_socket.recv(4096)
                    if not data:
                        self._log(f"客户端 {nickname} 断开连接", "ERROR")
                        break

                    # 更新最后活动时间
                    with self.lock:
                        if client_socket in self.clients:
                            self.clients[client_socket]["last_active"] = time.time()

                    # 检查是否为心跳包
                    if data == b"HEARTBEAT_REQUEST":
                        client_socket.send(b"HEARTBEAT_RESPONSE")
                        continue
                    # 聊天消息处理
                    try:
                        message = data.decode('utf-8')
                        self._broadcast(data, client_socket)
                        self._log(f"[{nickname}] {message}")
                    except UnicodeDecodeError:
                        self._log(f"收到无效消息格式", "ERROR")
                except socket.timeout:
                    continue
                except Exception as e:
                    self._log(f"处理消息时发生错误: {str(e)}", "ERROR")
                    break
        except Exception as e:
            self._log(f"客户端连接处理错误: {str(e)}", "ERROR")
        finally:
            self._cleanup_client(client_socket)

    def _handle_file_client(self, file_client_socket):
        """处理文件客户端连接（接收文件或支持下载）"""
        try:
            # 先接收一行，判断是上传还是下载
            header_buffer = b""
            while b"\n" not in header_buffer:
                chunk = file_client_socket.recv(1024)
                if not chunk:
                    self._log("文件客户端连接断开（未收到完整请求）", "ERROR")
                    return
                header_buffer += chunk
            header_line, remain = header_buffer.split(b"\n", 1)
            header = header_line.decode('utf-8')

            if header.startswith("FILE_INIT:"):
                # 上传流程（与之前一致）
                _, file_name, file_size, md5_hex, file_sender = header.split(":", 4)
                file_name = file_name.strip()
                file_size = int(file_size)
                md5_hex = md5_hex.strip()
                file_sender = file_sender.strip()
                if file_size <= 0 or file_size > 1024 * 1024 * 100:
                    self._log(f"无效的文件大小: {file_size} bytes", "ERROR")
                    return
                self._log(f"开始接收文件: {file_name} ({file_size} bytes) (MD5: {md5_hex})")
                received = len(remain)
                file_data = bytearray(remain)
                while received < file_size:
                    chunk = file_client_socket.recv(min(4096, file_size - received))
                    if not chunk:
                        self._log(f"文件接收中断: {file_name}", "ERROR")
                        return
                    file_data.extend(chunk)
                    received += len(chunk)
                    if received % (1024 * 1024) == 0 or received == file_size:
                        progress = received / file_size * 100
                        self._log(
                            f"接收 {file_name}: {progress:.1f}% ({received / 1024:.1f}KB/{file_size / 1024:.1f}KB)")
                received_md5 = hashlib.md5(file_data).hexdigest()
                if received_md5 != md5_hex:
                    self._log(f"文件校验失败: {file_name} (预期:{md5_hex} 实际:{received_md5})", "ERROR")
                    return
                save_dir = "server_files"
                os.makedirs(save_dir, exist_ok=True)
                save_path = os.path.join(save_dir, file_name)
                if os.path.exists(save_path):
                    name, ext = os.path.splitext(file_name)
                    timestamp = time.strftime("%Y%m%d_%H%M%S")
                    file_name = f"{name}_{timestamp}{ext}"
                    save_path = os.path.join(save_dir, file_name)
                with open(save_path, "wb") as f:
                    f.write(file_data)
                self._log(f"文件保存成功: {save_path} (MD5: {received_md5})")
                try:
                    client_ip, client_port = file_client_socket.getpeername()
                    sender_ip = f"{client_ip}:{client_port}"
                except:
                    sender_ip = "未知"
                notice = json.dumps({
                    "type": "file_notice",
                    "file_name": file_name,
                    "file_size": file_size,
                    "md5": md5_hex,
                    "sender": file_sender,
                    "ip": sender_ip
                }).encode('utf-8')
                self._broadcast(notice, None)
            elif header.startswith("GET:"):
                # 下载流程
                file_name = header[4:].strip()
                save_dir = "server_files"
                file_path = os.path.join(save_dir, file_name)
                if not os.path.exists(file_path):
                    file_client_socket.sendall(b"ERROR:FileNotFound\n")
                    self._log(f"客户端下载文件不存在: {file_name}", "ERROR")
                    return
                file_size = os.path.getsize(file_path)
                with open(file_path, "rb") as f:
                    file_data = f.read()
                md5_hex = hashlib.md5(file_data).hexdigest()
                # 发送文件头
                header = f"FILE_INIT:{file_name}:{file_size}:{md5_hex}\n".encode('utf-8')
                file_client_socket.sendall(header)
                # 发送文件内容
                sent = 0
                while sent < file_size:
                    chunk = file_data[sent:sent + 4096]
                    file_client_socket.sendall(chunk)
                    sent += len(chunk)
                self._log(f"客户端下载完成: {file_name} ({file_size} bytes)")
            else:
                self._log(f"收到无效文件端口请求: {header}", "ERROR")
        except Exception as e:
            self._log(f"文件客户端处理错误: {str(e)}", "ERROR")
        finally:
            try:
                file_client_socket.close()
            except:
                pass

    def _cleanup_client(self, client_socket):
        """清理客户端资源"""
        if client_socket in self.clients:
            info = self.clients[client_socket]
            try:
                # 获取客户端地址信息（备选方案）
                client_ip, client_port = client_socket.getpeername()
                address = f"{client_ip}:{client_port}"
            except:
                # 如果socket已关闭，使用存储的地址
                address = info.get('ip', '未知地址')
            del self.clients[client_socket]
            self._update_user_list()
            self._log(f"用户离线: '{info['nickname']}'  [{address}]", "ERROR")
        try:
            client_socket.close()
        except:
            pass

    def _update_user_list(self):
        """更新在线用户列表"""
        self.user_list.config(state="normal")
        self.user_list.delete(1.0, "end")
        users = [f"{info['nickname']} \n[{info['ip']}]\n" for info in self.clients.values()]
        self.user_list.insert("end", "\n".join(users))
        self.user_list.config(state="disabled")

    def _log(self, message, level="INFO"):
        """记录日志"""
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        color = "#333333"
        if level == "ERROR":
            color = "#FF0000"
        elif level == "WARNING":
            color = "#FF9900"
        elif level == "Login":
            color = "#009933"

        self.log_area.config(state="normal")
        self.log_area.insert("end", f"[{timestamp}] {message}\n", level)
        self.log_area.tag_config(level, foreground=color)
        self.log_area.config(state="disabled")
        self.log_area.see("end")

    def export_log(self):
        try:
            # 获取日志内容
            log_content = self.log_area.get("1.0", "end-1c")
            if not log_content.strip():
                messagebox.showwarning("提示", "日志内容为空")
                return

            # 弹出文件保存对话框
            file_path = filedialog.asksaveasfilename(
                defaultextension=".log",
                filetypes=[("日志文件", "*.log"), ("文本文件", "*.txt"), ("All Files", "*.*")],
                initialfile=f"chatlog_{time.strftime('%Y%m%d%H%M%S')}"
            )

            if not file_path:  # 用户取消保存
                return

            # 写入文件
            with open(file_path, "w", encoding="utf-8") as f:
                f.write(log_content)

            self._log(f"日志已导出至: {file_path}")
            messagebox.showinfo("成功", "日志导出完成")

        except Exception as e:
            self._log(f"日志导出失败: {str(e)}", "ERROR")
            messagebox.showerror("错误", f"导出失败: {str(e)}")

    def log_area_delete(self):
        self.log_area.config(state="normal")
        self.log_area.delete(1.0, "end")
        self.log_area.config(state="disabled")

    def run(self):
        self.window.mainloop()

    def on_closing(self):
        if messagebox.askokcancel("退出", "确定要退出吗？"):
            self.stop_server()
            self.window.destroy()


if __name__ == "__main__":
    server = ChatServer()
    server.run()
